from sklearn.linear_model import Ridge
import h5py
import pandas as pd
from matplotlib import pyplot as plt

f = h5py.File('out/dataset_order_cancel.h5', 'r')
ts = f['ts'][...]
x = f['x'][...]
y2 = f['y'][...]

mask = (y2 >= 0)
ts = ts[mask]
x = x[mask]
y2 = y2[mask]

print(x.shape)
print(y2.shape)

a = Ridge(alpha=0.1)
a.fit(x, y2)
s = a.score(x, y2)

# save
f = h5py.File('out/latency_model3.h5', 'w')
f['weights'] = a.coef_
f['bias'] = a.intercept_

ans = y2
pred = a.predict(x)

# yy plot
plt.scatter(pred, ans, label='yy', s=0.5, alpha=0.5)
# qq plot
plt.plot(sorted(pred), sorted(ans), label='qq', alpha=0.5)
# y=x
plt.plot([0.0, 4e9], [0.0, 4e9], alpha=0.5)

# time vs y
#plt.scatter(ts, pred, label='pred', alpha=0.5, s=0.5)
#plt.scatter(ts, ans, label='ans', alpha=0.5, color='orange', s=0.5)
#plt.ylim(0, 3e9)

plt.savefig('out/plot3.png', dpi=300)
print('score', s)
